package cn.lishiyuan.algorithm.tree;

import java.util.Objects;

public class RedBlackTree<T extends Comparable<T>> implements LeeTree<T>{

    private int size;

    private Node<T> root;

    private static final boolean RED = true;
    private static final boolean BLACK = false;

    public static class Node<T>{
        private T data;
        boolean color = RED;
        private Node<T> parent;
        private Node<T> left;
        private Node<T> right;
    }


    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public boolean contains(T element) {
        if(element == null){
            return false;
        }
        Node<T> node = findNode(element);
        return Objects.nonNull(node);
    }

    protected Node<T> findNode(T data){
        if(isEmpty()){
            return null;
        }
        Node<T> node = root;
        while (node!=null){
            if (data.compareTo(node.data) == 0){
                return node;
            }else if (data.compareTo(node.data) < 0){
                node = node.left;
            }else {
                node = node.right;
            }
        }
        return null;
    }

    @Override
    public T find(T element) {
        if(element == null){
            return null;
        }
        Node<T> node = findNode(element);
        return Objects.nonNull(node) ? node.data : null;
    }

    @Override
    public void add(T element) {
        if(element == null){
            throw new IllegalArgumentException("element is null");
        }
        // 新节点都是红色的
        Node<T> node = newNode(element);
        // 如果插入的节点是根节点，那我们直接改变它的颜色，把它变成黑色就可以了。
        if(root == null){
            root = node;
            root.color = BLACK;
            size++;
            return;
        }
        // 查找插入位置
        Node<T> parent = root;

        // 插入
        for(;true;){
            if (parent.data.compareTo(node.data)==0) {
                // 存在相同节点直接返回
                return;
            }else if (parent.data.compareTo(node.data) < 0){
                // 插入节点比当前节点大，向右查找
                if(parent.right == null){
                    // 直接插入
                    parent.right = node;
                    node.parent = parent;
                    break;
                }else {
                    parent = parent.right;
                }
            }else {
                // 插入节点比当前节点小，向右查找
                if(parent.left == null){
                    parent.left = node;
                    node.parent = parent;
                    break;
                }else {
                    parent = parent.left;
                }
            }
        }

        // 已经插入
        size++;

        Node<T> focusNode = node;

        // from 算法导论
        // 插入后修正
        // parent 和 node 都是红色（连续红节点检查，root元素不用检查）
        // 注意：节点为null被当作黑色节点来处理
        // 注意：我们的关注节点始终是红色节点
        // 注意：叔节点在左分支和右分支的处理是对称的
        for (; root != focusNode && focusNode.parent.color == RED;){

            Node<T> focusNodeParent = focusNode.parent;
            Node<T> focusNodeParentParent = focusNodeParent.parent;

            if(focusNodeParent == focusNodeParentParent.left){
                Node<T> uncle = focusNodeParentParent.right;
                // 叔叔节点是右节点
                if( colorOf(uncle) == RED){
                    // 叔叔节点是红色节点
                    /**
                     * CASE 1：如果关注节点是 a，它的叔叔节点 d 是红色
                     * 将关注节点 a 的父节点 b、叔叔节点 d 的颜色都设置成黑色；
                     * 将关注节点 a 的祖父节点 c 的颜色设置成红色；
                     * 关注节点变成 a 的祖父节点 c；
                     * 跳到 CASE 2 或者 CASE 3。
                     */
                    // 将关注节点 a 的父节点 b、叔叔节点 d 的颜色都设置成黑色；
                    setColor(uncle,BLACK);
                    setColor(focusNodeParent,BLACK);
                    // 将关注节点 a 的祖父节点 c 的颜色设置成红色
                    setColor(focusNodeParentParent,RED);
                    // 关注节点变成 a 的祖父节点 c
                    focusNode = focusNodeParentParent;
                }else {
                    // 叔叔节点是黑色节点
                    if(focusNodeParent.right == focusNode){
                        /**
                         * CASE 2：如果关注节点是 a，它的叔叔节点 d 是黑色，关注节点 a 是其父节点 b 的右子节点
                         *
                         * - 关注节点变成节点 a 的父节点 b；
                         * - 围绕新的关注节点b 左旋；
                         * - 跳到 CASE 3。
                         */
                        focusNode = focusNodeParent;
                        rotateLeft(focusNode);
                    }

                    /**
                     * CASE 3：如果关注节点是 a，它的叔叔节点 d 是黑色，关注节点 a 是其父节点 b 的左子节点
                     * - 围绕关注节点 a 的祖父节点 c 右旋；
                     * - 将关注节点 a 的父节点 b、兄弟节点 c 的颜色互换。
                     * - 调整结束。
                     */
                    setColor(focusNode.parent,BLACK);
                    setColor(focusNode.parent.parent,RED);
                    // 围绕关注节点 a 的祖父节点 c 右旋；
                    rotateRight(focusNode.parent.parent);

                }

            }else {
                Node<T> uncle = focusNodeParentParent.left;
                // 叔叔节点是左节点
                if(colorOf(uncle) == RED){
                    // 叔叔节点是红色节点
                    // 将关注节点 a 的父节点 b、叔叔节点 d 的颜色都设置成黑色；
                    setColor(uncle,BLACK);
                    setColor(focusNodeParent,BLACK);
                    // 将关注节点 a 的祖父节点 c 的颜色设置成红色
                    setColor(focusNodeParentParent,RED);
                    // 关注节点变成 a 的祖父节点 c
                    focusNode = focusNodeParentParent;
                }else {
                    // 叔叔节点是黑色节点
                    // 如果节点
                    if(focusNodeParent.left == focusNode){
                        focusNode = focusNodeParent;
                        rotateRight(focusNode);
                    }

                    setColor(focusNode.parent,BLACK);
                    setColor(focusNode.parent.parent,RED);

                    // 围绕关注节点 a 的祖父节点 c 右旋；
                    rotateLeft(focusNode.parent.parent);

                }
            }
        }

        root.color = BLACK;

    }

    @Override
    public T remove(T element) {
        if (element == null){
            throw new NullPointerException();
        }

        Node<T> node = findNode(element);

        if(node == null){
            // 未找到
            return null;
        }
        T returnData = node.data;

        // 有两个节点，转换成删除右子树最小节点操作
        if(node.left != null && node.right !=null){
            // 右树的的最小节点
            Node<T> min = node.right;
            while (min.left != null){
                min = min.left;
            }
            // 交换到当前位置
            node.data =  min.data;
            // 变成删除min
            node = min;
        }

        // 获取替换节点，上面已经处理了有两个节点的情况
        // 处理有一个子节点的情况
        Node<T> replaceNode = node.left != null ? node.left : node.right;

        if(replaceNode != null){
            replaceNode.parent = node.parent;
            if (node.parent == null)
                root = replaceNode;
            else if (node == node.parent.left)
                node.parent.left  = replaceNode;
            else
                node.parent.right = replaceNode;

            node.left = node.right = node.parent = null;

            // replaceNode只可能是红节点，进行染黑色处理
            /**
             * 首先是待删除黑节点有一颗子树的情况下，其子树必然是一个红节点，此时只需要将该红节点替换到待删除黑节点的位置即可，由于黑节点被删除，黑高度降低，那么此时我们就需要做出调整，将刚刚的红节点染色成黑节点
             *
             * ```
             * [p-any]      或者：[p-any]
             *    / \              / \
             * [any][黑]        [any][黑]
             *       /\               /\
             *   [红][nil]         [nil][红]
             * ```
             */
            if(node.color == BLACK){
                // replace节点必然为红色，需要改色
                replaceNode.color = BLACK;
            }
        }else if(node.parent == null) {
            // 左右节点都是空节点，当前节点是根节点的情况,直接root置空
            root = null;
        }else {
            // 左右子节点都是空节点，且当前节点不是根节点,有两种情况
            // 第一种情况是，原本节点就是没有左右节点的节点，如果该节点是一个黑节点，删除则破坏了黑高度，需要修复
            // 第二种情况是，原本节点有两个节点，而原本节点的右子树最小节点没有节点，被替换上来之后,需要删除,删除后破坏了黑高度,需要修复
            if(node.color == BLACK){
                fixAfterDeleted(node);
            }

            // 删除节点
            if(node.parent != null){
                // 清除节点
                if (node.parent.left == node){
                    node.parent.left = null;
                }else {
                    node.parent.right = null;
                }
            }
            node.left = node.right = node.parent = null;

        }
        size--;
        return returnData;
    }


    /**
     * 因为删除黑节点会改变黑高度，也会导致两个红节点相连
     * <br>
     * 又因为我们将删除节点的情况全部转换成了只有一个分支的情况，或者没有分支的情况。
     * <br>
     * 于是我们只需要分类讨论就行
     *
     * @param node
     */
    private void fixAfterDeleted(Node<T> node){
        while (node != root && colorOf(node) == BLACK){
            if(node.parent.right == node){
                // 当前黑节点在右分支
                Node<T> bro = node.parent.left;
                if(colorOf(bro) == RED){
                    /**
                     *
                     * 兄弟节点是红节点
                     * 此时父节点必然是黑节点，将其染红；将兄弟染黑
                     * 进行旋转，将黑兄提到父节点位，此时原本的兄弟分支的黑高没有变化，
                     * 原本兄弟分支的某个黑节点成了关注节点的新兄弟需要接着处理
                     *
                     *     [p]
                     *     / \
                     *[bro-红][node]
                     *   / \
                     * [黑][黑]
                     */
                    //当前兄弟节点必然是红色，父节点必然是黑色  父兄变色
                    setColor(bro,BLACK);
                    setColor(bro.parent,RED);
                    rotateRight(bro.parent);
                    // 此时关注节点的兄弟节点是原本bro的子节点
                }else {
                    if( Objects.isNull(bro) || (colorOf(bro.right) == BLACK && colorOf(bro.left) == BLACK) ){
                        /**
                         * 兄弟是黑节点，且右两个黑子节点（包括nil）
                         * 将兄弟节点变红，此时当前路径的左右两个节点都少了一个黑高度
                         * 将关注节点上移（此时，如果新的关注节点是红色，直接染黑，如果是黑色，那么就还需要继续操作）
                         *       [p]
                         *       / \
                         * [bro-黑][node]
                         *    /\
                         * [黑][黑]
                         */
                        setColor(bro,RED);
                        // 关注节点上移
                        node = node.parent;
                    }else{

                        // bro 同侧节点为红的情况
                        if(colorOf(bro.left) == RED){
                            /**
                             * bro是黑色，其同路子节点是红色
                             * 将该子节点设置为父节点颜色（必然是黑）
                             * 将bro的颜色设置为父节点颜色
                             * 将父节点设置为黑色
                             * 朝双黑节点旋转
                             * 完成调整
                             */
                            // bro的同路子节点变bro的颜色，bro的父节点变黑色
                            // 其实此时bro的同路子节点必然是红色，bro必然是黑色，bro的parent可黑可红
                            setColor(bro.left,colorOf(bro));
                            setColor(bro,colorOf(bro.parent));
                            setColor(bro.parent,BLACK);
                            rotateRight(node.parent);

                            node = root;
                        }else {
                            /**
                             * bro同侧子节点不为红，该异侧子节点设置为父节点一个颜色
                             * 将父节点设置为黑色，再将异侧子节点通过旋转提到bro位置
                             * 再围绕父节点进行旋转达到平衡
                             */
                            setColor(bro.right,colorOf(bro.parent));
                            setColor(bro.parent,BLACK);
                            rotateLeft(bro);
                            // 由于左旋，其右节点已经上来了
                            // 再围绕bro的parent的parent右旋
                            rotateRight(bro.parent.parent);
                            node = root;
                        }
                    }
                }

            }else {
                // 当前黑节点在左分支
                Node<T> bro = node.parent.right;
                if(colorOf(bro) == RED){
                    //当前兄弟节点必然是红色，父节点必然是黑色  父兄变色
                    setColor(bro,BLACK);
                    setColor(bro.parent,RED);
                    rotateLeft(bro.parent);
                    // 此时关注节点的兄弟节点是原本bro的子节点
                }else {
                    if( Objects.isNull(bro) || (colorOf(bro.right) == BLACK && colorOf(bro.left) == BLACK) ){
                        setColor(bro,RED);
                        // 关注节点上移
                        node = node.parent;
                    }else{

                        // bro 同侧节点为红的情况
                        if(colorOf(bro.right) == RED){
                            // bro的同路子节点变bro的颜色，bro的父节点变黑色
                            // 其实此时bro的同路子节点必然是红色，bro必然是黑色，bro的parent可黑可红
                            setColor(bro.right,colorOf(bro));
                            setColor(bro,colorOf(bro.parent));
                            setColor(bro.parent,BLACK);
                            rotateLeft(node.parent);

                            node = root;
                        }else {
                            setColor(bro.left,colorOf(bro.parent));
                            setColor(bro.parent,BLACK);
                            rotateRight(bro);
                            // 由于左旋，其右节点已经上来了
                            // 再围绕bro的parent的parent右旋
                            rotateLeft(bro.parent.parent);
                            node = root;
                        }
                    }
                }
            }
        }
        // 将关注节点变黑，会发生在关注节点为红色或者关注节点为root的情况下
        setColor(node,BLACK);

    }

    private boolean colorOf(Node<T> node){
        return node == null? BLACK : node.color;
    }

    private void setColor(Node<T> node , boolean color){
        if(node != null){
            node.color = color;
        }
    }


    /**
     * 左旋
     * @param node
     */
    protected void rotateLeft(Node<T> node) {
        // 将当前节点的右节点上提
        Node<T> newNode = node.right;
        // 将新的节点的左节点放到旧节点的右节点
        node.right = newNode.left;

        // 更新父节点
        if (node.right != null) {
            node.right.parent = node;
        }

        // 新节点的左节点设置为旧节点
        newNode.left = node;

        // 父节点处理
        Node<T> parent = node.parent;
        if (parent != null) {
            if(parent.left == node){
                parent.left = newNode;
            }else {
                parent.right = newNode;
            }
        }else {
            root = newNode;
        }

        // 新节点的父节点是旧节点的父节点
        newNode.parent = parent;
        // 旧节点的父节点就是新节点
        node.parent = newNode;
        // 将新节点设置到父节点的合适位置
    }

    /**
     * 右旋
     * @param node
     */
    protected void rotateRight(Node<T> node) {
        // 当前节点的左节点向上提
        Node<T> newNode = node.left;
        // 旧节点的左节点是新节点的右节点
        node.left = newNode.right;

        // 更新父节点
        if (node.left != null) {
            node.left.parent = node;
        }

        newNode.right = node;

        // parent
        Node<T> parent = node.parent;
        if (parent != null) {
            if(parent.left == node){
                parent.left = newNode;
            }else {
                parent.right = newNode;
            }
        }else {
            root = newNode;
        }
        newNode.parent = parent;
        node.parent = newNode;
    }

    @Override
    public String toString() {
        if(root == null){
            return "[]";
        }
        // 中序遍历
        String str = toString(root);
        return "["+str+"]";
    }

    private String toString(Node<T> node){
        if (node == null){
            return "";
        }
        StringBuilder builder = new StringBuilder();
        String leftStr = toString(node.left);
        if (!leftStr.isEmpty()){
            builder.append(leftStr).append(",");
        }
        String nodeStr = node.data.toString();
        builder.append(nodeStr);
        String rightStr = toString(node.right);
        if (!rightStr.isEmpty()){
            builder.append(",").append(rightStr);
        }
        return builder.toString();
    }

    private Node<T> newNode(T data){
        Node<T> tNode = new Node<>();
        tNode.color = RED;
        tNode.data = data;
        tNode.parent = null;
        tNode.left = null;
        tNode.right = null;
        return tNode;
    }
}
